package datastruct;

import java.util.Arrays;

public class UnionFind {


    // 父数组
    private final int[] p;
    //权重, 以 i 为根结点的子树的高度（引入了路径压缩以后该定义并不准确)
    private final int[] size;

    // 联通分量
    private int setCount;

    public UnionFind(int n) {
        this.p = new int[n];
        for (int i = 0; i < n; i++) {
            this.p[i] = i;
        }
        this.size = new int[n];
        Arrays.fill(size, 1);
        setCount = n;
    }

    /**
     * 找到元素 x 所在的集合的代表，该操作也可以用于判断两个元素是否位于同一个集合。
     *
     * @param x x
     * @return p[x]
     */
    public int find(int x) {
        if (p[x] != x) {
            p[x] = find(p[x]);
        }
        return p[x];
    }


    /**
     * 合并 x ,y
     *
     * @param x x
     * @param y y
     */
    public void union(int x, int y) {
        int pa = find(x);
        int pb = find(y);
        if (pa == pb) {
            return;
        }
        // 按权重 合并
        if (size[pa] > size[pb]) {
            p[pb] = p[pa];
            size[pa] += size[pb];
        } else {
            p[pa] = p[pb];
            size[pb] += size[pa];
        }
        --setCount;
    }


    /**
     * 判断是否联通
     *
     * @param x x
     * @param y y
     * @return bool true is connected
     */
    public boolean connected(int x, int y) {
        x = find(x);
        y = find(y);
        return x == y;
    }

}
